#include "KeyLock.h"

///Copyright (c) 2009-2010 T. Klepzig 

map<HWND, KeyLock*> KeyLock::instances;


KeyLock::KeyLock(HINSTANCE hInstance)
{
	this->hInstance = hInstance;
	this->locked = false;
	this->unauthorizedLoginOccured = false;

	this->primaryKeyVolumeLetter = 0;	
	this->secondaryKeyVolumeLetter = 0;	

	primaryKeySerial = Settings::GetData()->KeyLockData.PrimarySerialNumber;
	secondaryKeySerial = Settings::GetData()->KeyLockData.SecondarySerialNumber;
	lockAtUnplug = Settings::GetData()->KeyLockData.LockAtUnplug;
	wasLocked = Settings::GetData()->KeyLockData.Locked;

	keyLockHiddenWindow = new Window();
	keyLockHiddenWindow->Create(hInstance, TEXT("KeyLockHiddenWndClass"), HiddenWindowWndProc, 0, WS_POPUP, CS_NOCLOSE);
	instances[keyLockHiddenWindow->GetHWnd()] = this;
	

	lockWindow = new Window();
	lockWindow->SetTag(this);
	lockWindow->Create(hInstance,
					  TEXT("LockWindow"),
					  LockWindowWndProc,
					  WS_EX_TOOLWINDOW | WS_EX_TOPMOST,
					  WS_POPUP,
					  CS_NOCLOSE,
					  GetStockBrush(BLACK_BRUSH),
					  0,
					  LoadIcon(hInstance,MAKEINTRESOURCE(IDI_ICON1)));

	lockWindow->SetLocation(GetSystemMetrics(SM_XVIRTUALSCREEN), GetSystemMetrics(SM_YVIRTUALSCREEN));
	lockWindow->SetSize(GetSystemMetrics(SM_CXVIRTUALSCREEN), GetSystemMetrics(SM_CYVIRTUALSCREEN));

	registerDriveWindow = new RegisterDriveWindow(this);
	registerDriveWindow->Create(hInstance, LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1)));
	registerDriveWindow->DrivePluggedIn(DrivePluggedIn);

	WTSRegisterSessionNotification(keyLockHiddenWindow->GetHWnd(), NOTIFY_FOR_THIS_SESSION);	
}

KeyLock::~KeyLock()
{
	Settings::GetData()->KeyLockData.PrimarySerialNumber = primaryKeySerial;
	Settings::GetData()->KeyLockData.SecondarySerialNumber = secondaryKeySerial;
	Settings::GetData()->KeyLockData.LockAtUnplug = lockAtUnplug;

	Deactivate();

	instances.erase(keyLockHiddenWindow->GetHWnd());

	WTSUnRegisterSessionNotification(keyLockHiddenWindow->GetHWnd());

	delete lockWindow;
	lockWindow = 0;

	delete registerDriveWindow;
	registerDriveWindow = 0;

	delete keyLockHiddenWindow;
	keyLockHiddenWindow = 0;
}

CCState KeyLock::Activate()
{
	CCState result;
	if ((result = CassiopeiaComponent::Activate()) != CCState::Successful)
		return result;
		
	if (wasLocked || (lockAtUnplug && !IsKeyPluggedIn()))
		Lock();
	
    RegisterHotKey(keyLockHiddenWindow->GetHWnd(), 1, MOD_WIN, 75);

	return CCState::Successful;
}

CCState KeyLock::Deactivate()
{
	CCState result = CCState::NotReady;
	if ((result = CassiopeiaComponent::Deactivate()) != CCState::Successful || locked)
		return result;

    UnregisterHotKey(keyLockHiddenWindow->GetHWnd(), 1);	

	return CCState::Successful;
}


void KeyLock::SetPrimaryDevice()
{
	registerDriveWindow->SetTag((LPVOID)1);
	registerDriveWindow->Show();
}

void KeyLock::SetSecondaryDevice()
{
	registerDriveWindow->SetTag((LPVOID)2);
	registerDriveWindow->Show();
}

void KeyLock::ClearPrimaryDevice()
{
	primaryKeySerial = 0;
}

void KeyLock::ClearSecondaryDevice()
{
	secondaryKeySerial = 0;
}


bool KeyLock::GetLockAtUnplug()
{
	return lockAtUnplug;
}

void KeyLock::SetLockAtUnplug(bool value)
{
	lockAtUnplug = value;	

	if (active && lockAtUnplug)
	{
		if (!IsKeyPluggedIn())
			Lock();
	}

	Settings::GetData()->KeyLockData.LockAtUnplug = lockAtUnplug;
	Settings::Save();
}


bool KeyLock::IsKeyPluggedIn()
{
	bool result = false;
	set<TCHAR> drives = Utilities::GetDrives();

	set<TCHAR>::iterator iter;
	for (iter = drives.begin(); iter != drives.end(); iter++)
	{
		DWORD serial = Utilities::GetVolumeSerial(*iter);

		if (serial)
		{
			if (serial == primaryKeySerial)
			{
				primaryKeyVolumeLetter = *iter;
				result = true;
			}
			else if (serial == secondaryKeySerial)
			{
				secondaryKeyVolumeLetter = *iter;
				result = true;
			}
		}
	}

	return result;
}


void KeyLock::Lock()
{
	if (!active)
		return;

	if (locked)
		return;

	if (!primaryKeySerial && !secondaryKeySerial)
		return;
		
	locked = true;

	Settings::GetData()->KeyLockData.Locked = true;
	Settings::Save();

	LockWorkStation();
}

void KeyLock::Unlock()
{
	if (!locked)
		return;

	locked = false;

	Settings::GetData()->KeyLockData.Locked = false;
	Settings::Save();

	//poweron display
	PostMessage(keyLockHiddenWindow->GetHWnd(), WM_SYSCOMMAND, SC_MONITORPOWER, -1);

	unauthorizedLoginOccured = false;	
}


LRESULT CALLBACK KeyLock::HiddenWindowWndProc(LPVOID sender, HWND hWnd,UINT uMsg,WPARAM wParam,LPARAM lParam)
{	
	switch (uMsg)
	{
	case WM_WTSSESSION_CHANGE:
		{
			if (wParam == WTS_SESSION_UNLOCK)
			{
				if (instances[hWnd]->locked)
				{
					//unauthorized login
					instances[hWnd]->unauthorizedLoginOccured = true;

					//show lockWindows
					/*if (instances[hWnd]->lockWindows.empty())
					EnumDisplayMonitors(0, 0, EnumDisplayMonitorsProc, (LPARAM)instances[hWnd]);
					*/
					instances[hWnd]->lockWindow->SetLocation(GetSystemMetrics(SM_XVIRTUALSCREEN), GetSystemMetrics(SM_YVIRTUALSCREEN));
					instances[hWnd]->lockWindow->SetSize(GetSystemMetrics(SM_CXVIRTUALSCREEN), GetSystemMetrics(SM_CYVIRTUALSCREEN));
					instances[hWnd]->lockWindow->Show();
					instances[hWnd]->lockWindow->SetTopMost(true);

					LockWorkStation();
					//ToDo: maybe play an alarm sound
					//ToDo: start the webcam program
				}
			}
			if (wParam == WTS_SESSION_LOCK)
			{
				if (instances[hWnd]->locked)
				{
					//ToDo: set it as an option if the screen should be go down
					if (!instances[hWnd]->unauthorizedLoginOccured)
					{
						//only shutdown the display when the workstation is locked the first time
						//PostMessage(hWnd, WM_SYSCOMMAND, SC_MONITORPOWER, 2);
					}

					//start the screensaver
					PostMessage(hWnd, WM_SYSCOMMAND, SC_SCREENSAVE, 0);
				}
			}

			break;
		}
	case WM_DEVICECHANGE:
		{
			if (wParam == DBT_DEVICEARRIVAL || wParam == DBT_DEVICEREMOVECOMPLETE)
			{
				DEV_BROADCAST_HDR *hdr = (DEV_BROADCAST_HDR*)lParam;

				if (hdr->dbch_devicetype == DBT_DEVTYP_VOLUME)
				{
					DEV_BROADCAST_VOLUME *volume = (DEV_BROADCAST_VOLUME*)lParam;

					if (wParam == DBT_DEVICEARRIVAL)
					{	
						TCHAR volumeLetter = Utilities::GetVolumeLetter(volume->dbcv_unitmask);
						DWORD serial = Utilities::GetVolumeSerial(volumeLetter);

						if (serial == instances[hWnd]->primaryKeySerial)
						{
							instances[hWnd]->primaryKeyVolumeLetter = volumeLetter;
							instances[hWnd]->Unlock();
						}
						else if (serial == instances[hWnd]->secondaryKeySerial)
						{
							instances[hWnd]->secondaryKeyVolumeLetter = volumeLetter;
							instances[hWnd]->Unlock();
						}
					}
					else if (wParam == DBT_DEVICEREMOVECOMPLETE)
					{		
						if (instances[hWnd]->lockAtUnplug && !instances[hWnd]->locked)
						{
							TCHAR volumeLetter = Utilities::GetVolumeLetter(volume->dbcv_unitmask);

							if (!instances[hWnd]->IsKeyPluggedIn())
							{
								if (volumeLetter == instances[hWnd]->primaryKeyVolumeLetter || volumeLetter == instances[hWnd]->secondaryKeyVolumeLetter)
									instances[hWnd]->Lock();
							}
						}
					}
				}
			}

			break;
		}
	case WM_HOTKEY:
		{
			if (wParam == 1)
			{
				instances[hWnd]->Lock();
			}

			break;
		}
	}

	return DefWindowProc(hWnd,uMsg,wParam,lParam);
}

void KeyLock::DrivePluggedIn(RegisterDriveWindow* sender, TCHAR driveLetter, DWORD driveSerial, LPVOID parent, LPVOID tag)
{
	if ((int)tag == 1) //primary device
	{
		((KeyLock*)parent)->primaryKeySerial = driveSerial;
		Settings::GetData()->KeyLockData.PrimarySerialNumber = driveSerial;
		Settings::Save();
	}
	else if ((int)tag == 2) //secondary device
	{
		((KeyLock*)parent)->secondaryKeySerial = driveSerial;
		Settings::GetData()->KeyLockData.SecondarySerialNumber = driveSerial;
		Settings::Save();
	}

}

LRESULT CALLBACK KeyLock::LockWindowWndProc(LPVOID sender, HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{	
	Window* window = (Window*)sender;
	KeyLock* instance = (KeyLock*)((Window*)sender)->GetTag();

	switch(uMsg)
	{
	case WM_LBUTTONUP:
		{
			if (!instance->locked)
			{
				/*list<Window*>::iterator iter;
				for(iter = instance->lockWindows.begin(); iter != instance->lockWindows.end(); iter++)
				{
				(*iter)->Close();

				delete *iter;
				*iter = 0;
				}

				instance->lockWindows.clear();*/

				window->Hide();
			}

			break;
		}
	case WM_PAINT:
		{
			RECT clientRect;
			GetClientRect(hWnd, &clientRect);

			HDC          hdc;
			PAINTSTRUCT  ps;

			hdc = BeginPaint(hWnd, &ps);

			Graphics g(hdc);

			//SolidBrush brushRed(Color(166, 16, 16));
			SolidBrush brushBlue(Color(19, 57, 89));

			//g.FillPie(&brushRed, clientRect.left + 1, clientRect.top, 20, 20, 90, 180);
			//g.FillPie(&brushRed, clientRect.right - 22, clientRect.top, 20, 20, -90, 180);
			g.FillRectangle(&brushBlue, clientRect.left, clientRect.top, RECTWIDTH(clientRect), 20); 

			//g.FillPie(&brushRed, clientRect.left + 1, clientRect.bottom - 20, 20, 20, 90, 180);
			//g.FillPie(&brushRed, clientRect.right - 22, clientRect.bottom - 20, 20, 20, -90, 180);
			g.FillRectangle(&brushBlue, clientRect.left, clientRect.bottom - 20, RECTWIDTH(clientRect), 20); 

			Utilities::DrawDefaultText(clientRect, hdc, TEXT("Unauthorized Access"), 22, RGB(255, 255, 255), DT_LEFT | DT_TOP | DT_SINGLELINE, FW_BOLD, RECTWIDTH(clientRect) / 14, 25);
			Utilities::DrawDefaultText(clientRect, hdc, TEXT("Unauthorized Access"), 22, RGB(255, 255, 255), DT_RIGHT | DT_BOTTOM | DT_SINGLELINE, FW_BOLD, -(RECTWIDTH(clientRect) / 14), -25);

			EndPaint(hWnd, &ps);

			break;
		}
	}

	return DefWindowProc(hWnd,uMsg,wParam,lParam);
}

//BOOL CALLBACK KeyLock::EnumDisplayMonitorsProc(HMONITOR hMonitor, HDC hdc, LPRECT lpRect, LPARAM lParam)
//{
//	KeyLock* instance = (KeyLock*)lParam;
//
//	Window* lockWindow = new Window();
//	lockWindow->SetTag(instance);
//	lockWindow->Create(instance->hInstance,
//					  TEXT("LockWindow"),
//					  LockWindowWndProc,
//					  WS_EX_TOOLWINDOW | WS_EX_TOPMOST,
//					  WS_POPUP,
//					  CS_NOCLOSE,
//					  GetStockBrush(BLACK_BRUSH),
//					  0,
//					  LoadIcon(instance->hInstance,MAKEINTRESOURCE(IDI_ICON1)));
//
//	lockWindow->SetLocation(lpRect->left, lpRect->top);
//	lockWindow->SetSize(PRECTWIDTH(lpRect), PRECTHEIGHT(lpRect));
//	lockWindow->Show();
//	lockWindow->SetTopMost(true);
//
//	instance->lockWindows.push_back(lockWindow);
//
//	return TRUE;
//}
